package com.base.service;



import com.base.bean.EsException;
import com.base.bean.EsHistogramAggregationReqPO;
import com.base.bean.EsMultiAggregationReqPO;
import com.base.bean.aggregations.*;
import com.google.common.collect.Lists;
import com.sun.istack.internal.NotNull;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.histogram.Histogram;
import org.elasticsearch.search.aggregations.bucket.histogram.HistogramAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.histogram.ParsedHistogram;
import org.elasticsearch.search.aggregations.bucket.terms.IncludeExclude;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.Cardinality;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.springframework.stereotype.Repository;
import org.springframework.util.CollectionUtils;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 参考链接 https://blog.csdn.net/yzh_1346983557/article/details/122674983
 */
@Repository
@Slf4j
public class EsAggrSearchRepository {

    @Resource
    private RestHighLevelClient highLevelClient;


    /**
     * 聚合统计
     *
     * @param reqPO 查询请求对象
     *
     * @return 聚合统计结果
     */
    public AggregationRespPO searchAggregation(EsAggregationReqPO reqPO) {
        // 封装查询源对象
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        // 查询条件
        sourceBuilder.query(reqPO.getQuery());

        // 排序字段
        if (StringUtils.isNotBlank(reqPO.getSortField()) && reqPO.getSort() != null) {
            FieldSortBuilder order = new FieldSortBuilder(reqPO.getSortField()).order(reqPO.getSort());
            sourceBuilder.sort(order);
        }

        // 页大小0，只返回聚合结果
        sourceBuilder.size(0);

        // 聚合分桶。创建terms桶聚合，聚合名字=terms_by_XXX, 字段=XXX
        TermsAggregationBuilder termsAggregationBuilder = AggregationBuilders.terms("terms_by_" + reqPO.getField()).field(reqPO.getField());
        if (reqPO.getFieldSize() != null) {
            termsAggregationBuilder.size(reqPO.getFieldSize());
        }
        // 聚合分桶的字段值包含
        if (reqPO.getIncludeValues() != null) {
            termsAggregationBuilder.includeExclude(new IncludeExclude(reqPO.getIncludeValues(), new String[0]));
        }
        // 聚合分组条件
        sourceBuilder.aggregation(termsAggregationBuilder);

        // 设置索引、source
        SearchRequest searchRequest = new SearchRequest(reqPO.getIndex()).source(sourceBuilder);

        // 查询结果
        SearchResponse searchResponse;
        try {
            log.info("es聚合统计请求: index={}, 请求对象source={}", reqPO.getIndex(), searchRequest.source());
            searchResponse = highLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            log.info("es聚合统计响应结果：{}", searchResponse);
        } catch (IOException e) {
            log.error("es查询，IO异常！searchRequest={}", searchRequest, e);
            throw new EsException("es查询，IO异常！");
        }

        if (RestStatus.OK.equals(searchResponse.status())) {
            long totalHits = searchResponse.getHits().getTotalHits().value;
            // 遍历terms聚合结果
            Terms terms = searchResponse.getAggregations().get(termsAggregationBuilder.getName());
            List<AggregationBucketPO> bucketPOList = convertTerms(terms);
            return new AggregationRespPO(totalHits, bucketPOList);
        } else {
            log.error("es查询返回的状态码异常！searchResponse.status={}, searchRequest={}", searchResponse.status(), searchRequest);
            throw new EsException("es查询返回的状态码异常！");
        }
    }

    /**
     * 二级聚合查询
     *
     * @param reqPO 查询请求对象
     *
     * @return 聚合查询结果
     */
    public List<AggregationBucketPO> searchMultiAggregation(EsMultiAggregationReqPO reqPO) {
        // 封装查询源对象
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        // 查询条件
        sourceBuilder.query(reqPO.getQuery());

        // 排序字段
        if (StringUtils.isNotBlank(reqPO.getSortField()) && reqPO.getSort() != null) {
            FieldSortBuilder order = new FieldSortBuilder(reqPO.getSortField()).order(reqPO.getSort());
            sourceBuilder.sort(order);
        }

        // 页大小0，只返回聚合结果
        sourceBuilder.size(0);

        // 聚合分桶。创建terms桶聚合，聚合名字=terms_by_XXX, 字段=XXX
        TermsAggregationBuilder termsAggregationBuilder = AggregationBuilders.terms("terms_by_" + reqPO.getField()).field(reqPO.getField());
        if (reqPO.getFieldSize() != null) {
            termsAggregationBuilder.size(reqPO.getFieldSize());
        }
        // 二级聚合分桶
        TermsAggregationBuilder subTermsAggregationBuilder = AggregationBuilders.terms("terms_by_" + reqPO.getSubField()).field(reqPO.getSubField());
        if (reqPO.getSubFieldSize() != null) {
            subTermsAggregationBuilder.size(reqPO.getSubFieldSize());
        }
        termsAggregationBuilder.subAggregation(subTermsAggregationBuilder);

        // 聚合分组条件
        sourceBuilder.aggregation(termsAggregationBuilder);

        // 设置索引、source
        SearchRequest searchRequest = new SearchRequest(reqPO.getIndex()).source(sourceBuilder);

        // 查询结果
        SearchResponse searchResponse = null;
        try {
            log.info("es二级聚合查询: index={}, 请求对象source={}", reqPO.getIndex(), searchRequest.source());
            // 执行搜索
            searchResponse = highLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            log.info("es二级聚合结果：searchResponse={}", searchResponse);
        } catch (IOException e) {
            log.error("es查询，IO异常！searchRequest={}", searchRequest, e);
            throw new EsException("es查询，IO异常！");
        }

        if (RestStatus.OK.equals(searchResponse.status())) {
            // 遍历terms聚合结果
            Terms terms = searchResponse.getAggregations().get(termsAggregationBuilder.getName());

            return terms.getBuckets().stream().map(bucket -> {
                // 一级聚合分桶的数据
                String key = bucket.getKeyAsString();
                long docCount = bucket.getDocCount();
                // 二级聚合分桶的数据
                Terms subTerms = bucket.getAggregations().get(subTermsAggregationBuilder.getName());
                List<AggregationBucketPO> subBucketList = convertTerms(subTerms);
                return new AggregationBucketPO(key, docCount, subBucketList);
            }).collect(Collectors.toList());
        } else {
            log.error("es查询返回的状态码异常！searchResponse.status={}, searchRequest={}", searchResponse.status(), searchRequest);
            throw new EsException("es查询返回的状态码异常！");
        }
    }

    /**
     * 直方图聚合统计
     *
     * @param reqPO 聚合统计对象
     *
     * @return 直方图聚合统计结果
     */
    public List<AggregationBucketPO> searchHistogramAggregation(EsHistogramAggregationReqPO reqPO) {
        // 封装查询源对象
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        // 查询条件
        sourceBuilder.query(reqPO.getQuery());

        // 排序字段
        if (StringUtils.isNotBlank(reqPO.getSortField()) && reqPO.getSort() != null) {
            FieldSortBuilder order = new FieldSortBuilder(reqPO.getSortField()).order(reqPO.getSort());
            sourceBuilder.sort(order);
        }

        // 页大小0，只返回聚合结果
        sourceBuilder.size(0);

        // 直方图聚合分桶。创建histogram桶聚合，聚合名字=histogram_by_XXX, 字段=XXX
        HistogramAggregationBuilder histogramAggregationBuilder = AggregationBuilders.histogram("histogram_by_" + reqPO.getField())
                .field(reqPO.getField()).interval(reqPO.getInterval())
                .extendedBounds(reqPO.getStartTime(), reqPO.getEndTime());
        // histogram 类型聚合的桶的key的公式：bucket_key = Math.floor((value - offset) / interval) * interval + offset
        // 当value为每一天的00:00:00时，由于有一个+8的时差的问题，value/interval不能整除，会导致bucket_key为每天的早上8点。interval=1天，offset+16小时
        if (reqPO.getInterval() == (1000 * 60 * 60 * 24)) {
            histogramAggregationBuilder.offset(1000 * 60 * 60 * 16);
        }

        // 聚合分组条件
        sourceBuilder.aggregation(histogramAggregationBuilder);

        // 设置索引、请求source
        SearchRequest searchRequest = new SearchRequest(reqPO.getIndex()).source(sourceBuilder);

        // 查询结果
        SearchResponse searchResponse;
        try {
            log.info("es直方图聚合统计请求: index={}, 请求对象source={}", reqPO.getIndex(), searchRequest.source());
            searchResponse = highLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            log.info("es直方图聚合统计响应：{}", searchResponse);
        } catch (IOException e) {
            log.error("es查询，IO异常！searchRequest={}", searchRequest, e);
            throw new EsException("es查询，IO异常！");
        }

        if (RestStatus.OK.equals(searchResponse.status())) {
            ParsedHistogram parsedHistogram = searchResponse.getAggregations().get(histogramAggregationBuilder.getName());
            List<? extends Histogram.Bucket> buckets = parsedHistogram.getBuckets();
            if (CollectionUtils.isEmpty(buckets)) {
                return Lists.newArrayList();
            }

            // 响应结果转换
            return buckets.stream().map(this::getBucketPO).collect(Collectors.toList());
        } else {
            log.error("es查询返回的状态码异常！searchResponse.status={}, searchRequest={}", searchResponse.status(), searchRequest);
            throw new EsException("es查询返回的状态码异常！");
        }
    }

    /**
     * 直方图二级聚合统计
     *
     * @param reqPO 聚合统计对象
     *
     * @return 直方图二级聚合统计结果
     */
    public List<AggregationBucketPO> searchHistogramMultiAggregation(EsHistogramAggregationReqPO reqPO) {
        // 封装查询源对象
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        // 查询条件
        sourceBuilder.query(reqPO.getQuery());

        // 排序字段
        if (StringUtils.isNotBlank(reqPO.getSortField()) && reqPO.getSort() != null) {
            FieldSortBuilder order = new FieldSortBuilder(reqPO.getSortField()).order(reqPO.getSort());
            sourceBuilder.sort(order);
        }

        // 页大小0，只返回聚合结果
        sourceBuilder.size(0);

        // 直方图聚合分桶。创建histogram桶聚合，聚合名字=histogram_by_XXX, 字段=XXX
        HistogramAggregationBuilder histogramAggregationBuilder = AggregationBuilders.histogram("histogram_by_" + reqPO.getField())
                .field(reqPO.getField()).interval(reqPO.getInterval())
                .extendedBounds(reqPO.getStartTime(), reqPO.getEndTime());
        // histogram 类型聚合的桶的key的公式：bucket_key = Math.floor((value - offset) / interval) * interval + offset
        // 当value为每一天的00:00:00时，由于有一个+8的时差的问题，value/interval不能整除，会导致bucket_key为每天的早上8点。interval=1天，offset+16小时
        if (reqPO.getInterval() == (1000 * 60 * 60 * 24)) {
            histogramAggregationBuilder.offset(1000 * 60 * 60 * 16);
        }

        // 二级聚合分桶
        TermsAggregationBuilder subTermsAggregationBuilder = AggregationBuilders.terms("terms_by_" + reqPO.getSubField()).field(reqPO.getSubField());
        if (reqPO.getSubFieldSize() != null) {
            subTermsAggregationBuilder.size(reqPO.getSubFieldSize());
        }
        if (reqPO.getSubIncludeValues() != null) {
            subTermsAggregationBuilder.includeExclude(new IncludeExclude(reqPO.getSubIncludeValues(), new String[0]));
        }
        histogramAggregationBuilder.subAggregation(subTermsAggregationBuilder);

        // 聚合分组条件
        sourceBuilder.aggregation(histogramAggregationBuilder);

        // 设置索引、请求source
        SearchRequest searchRequest = new SearchRequest(reqPO.getIndex()).source(sourceBuilder);

        // 查询结果
        SearchResponse searchResponse;
        try {
            log.info("es直方图二级聚合统计请求: index={}, 请求对象source={}", reqPO.getIndex(), searchRequest.source());
            // 执行搜索
            searchResponse = highLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            log.info("es直方图二级聚合统计响应：{}", searchResponse);
        } catch (IOException e) {
            log.error("es查询，IO异常！searchRequest={}", searchRequest, e);
            throw new EsException("es查询，IO异常！");
        }

        if (RestStatus.OK.equals(searchResponse.status())) {
            ParsedHistogram parsedHistogram = searchResponse.getAggregations().get(histogramAggregationBuilder.getName());
            List<? extends Histogram.Bucket> buckets = parsedHistogram.getBuckets();
            if (CollectionUtils.isEmpty(buckets)) {
                return Lists.newArrayList();
            }

            // 响应结果转换
            return buckets.stream().map(bucket -> getBucketPO(subTermsAggregationBuilder, bucket)).collect(Collectors.toList());
        } else {
            log.error("es查询返回的状态码异常！searchResponse.status={}, searchRequest={}", searchResponse.status(), searchRequest);
            throw new EsException("es查询返回的状态码异常！");
        }
    }


    /**
     * 聚合的分页查询
     *
     * @param queryPO 查询请求对象
     *
     * @return 聚合分页查询结果
     */
    public EsQueryRespPO<AggregationBucketPO> searchAggregation(EsQueryReqPO queryPO) {
        // 设置索引
        SearchRequest searchRequest = new SearchRequest(queryPO.getIndex());

        // 封装查询源对象
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        searchRequest.source(sourceBuilder);

        // 查询条件
        sourceBuilder.query(queryPO.getQuery());

        // 排序字段
        if (StringUtils.isNotBlank(queryPO.getSortField()) && queryPO.getSort() != null) {
            FieldSortBuilder order = new FieldSortBuilder(queryPO.getSortField()).order(queryPO.getSort());
            sourceBuilder.sort(order);
        }

        // 页大小0，只返回聚合结果
        sourceBuilder.size(0);

        // 设置聚合查询,可以设置多个聚合查询条件，只要聚合查询命名不同就行
        // 聚合分组条件, group by
        sourceBuilder.aggregation(queryPO.getTermsAggregation());
        // 聚合统计条件, count分组后的数据，计算分组后的总大小
        sourceBuilder.aggregation(queryPO.getCardinalityAggregation());

        // 查询结果
        SearchResponse searchResponse = null;
        try {
            // log.info("es查询请求对象：searchRequest={}", searchRequest);
            log.info("es查询请求对象source：sourceBuilder={}", searchRequest.source());
            // 执行搜索
            searchResponse = highLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            log.info("es查询响应结果：searchResponse={}", searchResponse);
        } catch (IOException e) {
            log.error("es查询，IO异常！searchRequest={}", searchRequest, e);
            return EsQueryRespPO.error("es查询，IO异常！");
        }

        if (RestStatus.OK.equals(searchResponse.status())) {
            // 解析对象
            Aggregations aggregations = searchResponse.getAggregations();

            long docTotal = searchResponse.getHits().getTotalHits().value;

            // 遍历terms聚合结果
            Terms terms = aggregations.get(queryPO.getTermsAggregation().getName());
            List<AggregationBucketPO> bucketList = terms.getBuckets().stream().map(bucket -> {
                String key = bucket.getKeyAsString();
                long docCount = bucket.getDocCount();
                return new AggregationBucketPO(key, docCount, docTotal);
            }).collect(Collectors.toList());

            // 总数量
            Cardinality cardinality = aggregations.get(queryPO.getCardinalityAggregation().getName());
            long totalHits = cardinality.getValue();

            return EsQueryRespPO.success(bucketList, queryPO.getPageNum(), queryPO.getPageSize(), totalHits);
        } else {
            log.error("es查询返回的状态码异常！searchResponse.status={}, searchRequest={}", searchResponse.status(), searchRequest);
            return EsQueryRespPO.error("es查询返回的状态码异常！");
        }
    }

    @NotNull
    private AggregationBucketPO getBucketPO(Histogram.Bucket bucket) {
        Double timestamp = (Double) bucket.getKey();
        String key = String.valueOf(timestamp.longValue());
        long docCount = bucket.getDocCount();
        return new AggregationBucketPO(key, docCount);
    }

    @NotNull
    private AggregationBucketPO getBucketPO(TermsAggregationBuilder subTermsAggregationBuilder, Histogram.Bucket bucket) {
        Double timestamp = (Double) bucket.getKey();
        String key = String.valueOf(timestamp.longValue());
        long docCount = bucket.getDocCount();
        Terms terms = bucket.getAggregations().get(subTermsAggregationBuilder.getName());
        List<AggregationBucketPO> bucketPOS = convertTerms(terms);
        return new AggregationBucketPO(key, docCount, bucketPOS);
    }

    private List<AggregationBucketPO> convertTerms(Terms terms) {
        if (CollectionUtils.isEmpty(terms.getBuckets())) {
            return Lists.newArrayList();
        }

        return terms.getBuckets().stream().map(bucket -> {
            String key = bucket.getKeyAsString();
            long docCount = bucket.getDocCount();
            return new AggregationBucketPO(key, docCount);
        }).collect(Collectors.toList());
    }
}
